package com.jboltai.capability.vdb.types.elasticsearch;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.json.jackson.JacksonJsonpMapper;
import co.elastic.clients.transport.ElasticsearchTransport;
import co.elastic.clients.transport.rest_client.RestClientTransport;
import org.apache.http.HttpHost;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.CredentialsProvider;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestClientBuilder;

public class Client {
    private RestClient restClient;

    private ElasticsearchTransport transport;

    private ElasticsearchClient esClient;

    public Client(ElasticSearchVDBResource resource) {


        RestClientBuilder builder = RestClient.builder(HttpHost.create(resource.getUrl()));

        RestClient restClient = null;

        if (resource.getUsername() != null && resource.getPassword() != null) {
            final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
            credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(resource.getUsername(), resource.getPassword()));
            builder.setHttpClientConfigCallback(b -> {
                        return b.setDefaultCredentialsProvider(credentialsProvider) ;
                    });

        }

        restClient = builder.build();


        ElasticsearchTransport transport = new RestClientTransport(restClient, new JacksonJsonpMapper());

        ElasticsearchClient client = new ElasticsearchClient(transport);

        this.restClient = restClient;
        this.transport = transport;
        this.esClient = client;
    }

    public RestClient getRestClient() {
        return restClient;
    }

    public void setRestClient(RestClient restClient) {
        this.restClient = restClient;
    }

    public ElasticsearchTransport getTransport() {
        return transport;
    }

    public void setTransport(ElasticsearchTransport transport) {
        this.transport = transport;
    }

    public ElasticsearchClient getEsClient() {
        return esClient;
    }

    public void setEsClient(ElasticsearchClient esClient) {
        this.esClient = esClient;
    }

    public void close() {
        try {
            transport.close();
            restClient.close();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
}
